#!/usr/bin/perl

# -*- mode: cperl; cperl-continued-brace-offset: -4; indent-tabs-mode: nil; -*-
# vim:shiftwidth=2:tabstop=8:expandtab:textwidth=78:softtabstop=4:ai:

#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# (C) Copyright Ticketmaster, Inc. 2007
#

# ============================================================================
#
#  $Id: onall 16 2009-03-23 23:24:53Z phil@ipom.com $
#
#  onall v2.11 -- A parallelized system-administration tool.
#    Written by Rafi Khardalian.
# 
#  Usage: onall [OPTIONS] <command>
#      onall [OPTIONS] --script <script> [<script_args>]
#      onall [OPTIONS] --copy <file> <dest_path>
# ============================================================================

use strict;

our $ssh = 'ssh';
our $ssh_add = 'ssh-add';
our $bash = 'bash';
our $fping = 'fping';

our %default = (
    user        => 'root',
    rate        => 20,
    timeout     => 300,
    delay       => 0,
    syslist     => '-',
);

# ============================================================================

use POSIX qw(:signal_h :errno_h :sys_wait_h);
use Socket;
use Getopt::Std;
use Term::ReadLine;
use File::Basename;
use Getopt::Long;
use Data::Dumper;


our $version = '2.11';
our $exec_cmd = '';
our $children = 0;
our %children = ();

$| = 1;
$SIG{CHLD} = sub { \&reaper };


our %opt;
Getopt::Long::Configure ('bundling');
Getopt::Long::GetOptions(
    \%opt,
    'copy|c=s',
    'delay|d=s',
    'help|h',
    'linenums|l',
    'logpath|o=s',
    'no_syntax_check|n',
    'ping|p',
    'quiet|q',
    'rate|r=s',
    'script|e=s',
    'ssh_arguments|A=s',
    'ssh_no_keys|k',
    'ssh_quiet|Q',
    'suppress|s=s',
    'syslist|f=s',
    'timeout|t=s',
    'unbuffered|b',
    'user|u=s',
    'yes|y',
    'printdots|x',
) or usage();

$opt{rate} = $default{rate}
    unless (exists $opt{rate} or $opt{rate} =~ m/\d+/);
$opt{delay} = $default{delay}
    unless (exists $opt{delay} or $opt{delay} =~ m/\d+/);
$opt{timeout} = $default{timeout}
    unless (exists $opt{timeout} or $opt{timeout} =~ m/\d+/);
$opt{user} = $default{user}
    unless (exists $opt{user});
$opt{syslist} = $default{syslist}
    unless (exists $opt{syslist});
delete $opt{linenums}
    if (exists $opt{suppress});

if (exists $opt{help}) { usage('extended'); exit 1; }

our $ssh_args = "-l $opt{user} -xT " 
                . '-o "PasswordAuthentication no" ' 
                . '-o "StrictHostKeyChecking no" ';

if (exists $opt{ssh_quiet}) {
    $ssh_args .= '-q ';
}

if (exists $opt{ssh_arguments}) {
    $ssh_args .= $opt{ssh_arguments};
}

if ( exists $opt{script} )
{
    $opt{script_args} = shift @ARGV || '';

    exit_error("ERROR: Specified script does not exist\n", 0)
    unless (-f $opt{script});

    $exec_cmd = 'TMPFILE="/tmp/onall_script.$$" ; cat - >$TMPFILE ; '
                . 'chmod +x $TMPFILE && $TMPFILE ' 
                . $opt{script_args} . '; rm $TMPFILE'; 
}
elsif ( exists $opt{copy} )
{
    $opt{dest} = shift @ARGV || $opt{copy};

    exit_error("ERROR: Specified file does not exist\n", 0)
        unless (-f $opt{copy});

    $exec_cmd = 'FILE="' . $opt{dest} . '" ; if [ -d $FILE ]; then '
                 . 'FILE="' . $opt{dest} . '/' . basename($opt{copy})
                 . '"; fi; cat - >$FILE ; chmod +x $FILE';
}
elsif (scalar(@ARGV) >= 1)
{
    $exec_cmd = join(" ", @ARGV);
}
else
{
    exit_error("ERROR: Nothing to do. Exiting.\n\n", 1);
}

exit_error("ERROR: Logpath $opt{logpath} is not a directory\n")
    if (exists $opt{logpath} and not -d $opt{logpath});

my @tools_to_check = ($ssh, $ssh_add, $bash);
push(@tools_to_check, $fping) if ($opt{ping});
@tools_to_check = grep({ $_ ne $ssh_add } @tools_to_check)
    if ($opt{'ssh_no_keys'});
@tools_to_check = grep({ $_ ne $bash } @tools_to_check)
    if ($opt{'no_syntax_check'});

check_tools(\@tools_to_check) or exit 1;
unless ($opt{'ssh_no_keys'}) {
    check_ssh_keys() or exit 1;
}
unless ($opt{'no_syntax_check'}) {
    check_bash($exec_cmd) or exit 1;
}

our $hostlist = build_hostlist($opt{syslist});
our $tcount = scalar(keys %{$hostlist->{reachable}});
our $ucount = scalar(keys %{$hostlist->{unreachable}});

exit_error("ERROR: All specified systems are unreachable.\n", 0)
    if ($tcount < 1);


unless (exists $opt{quiet})
{
    print STDERR "\n" . $tcount . " Target(s):\n\n";
    print_list(sort values %{$hostlist->{reachable}});

    if ($ucount >= 1)
    {
        print STDERR "\n" . $ucount . " Unreachable:\n\n";
        print_list(sort values %{$hostlist->{unreachable}});
    }

 
    print STDERR "\nThe following command will be executed:\n\t";
    if (exists $opt{script}) {
        print STDERR "Script: $opt{script}\n";
    } elsif (exists $opt{copy}) {
        print STDERR "Copy: $opt{copy} --> $opt{dest}\n";
    } else {
        print STDERR $exec_cmd, "\n";
    }

    if (not exists $opt{yes}) 
    {
        my $term = new Term::ReadLine '';
        my $prompt = '';

        while (not $prompt =~ m/^\s*[y|n]/i)
        {
            $prompt = $term->readline('Continue (Y/N):');
        }

        unless ($prompt =~ m/^y/i)
        {
            print "\nUser aborted. Exiting.\n";
            exit 1
        }
    }

    print "\n";
}


for my $machine (sort keys %{$hostlist->{reachable}})
{
    chomp $machine;

    while ($children >= $opt{rate}) 
    {
        &reaper;
        select(undef, undef, undef, 0.25);
    }

    $children++;    

    defined(my $pid = fork) or 
        die "Cannot fork: $!";

    $children{$pid} = 1;

    if (not $pid) 
    {
        new_ssh_child($machine);
    }

    &reaper;
}


while (wait > -1) {
    &reaper;
    select(undef, undef, undef, 0.25);
}

print "\n" unless (exists $opt{quiet});


# ============================================================================

sub new_ssh_child
{
    my $machine = shift;
    my ($hostname, @pids, @output);
    my $linecounter = 0;

    $hostname = $hostlist->{reachable}->{$machine};

    $SIG{ALRM} = sub 
    { 
        print STDERR "$hostname: Timeout exceeded \[$opt{timeout}\] -- ";
        print STDERR 'killing \(' . join(' ', @pids) . "\).\n";

        for (@pids) { kill KILL => $_; }
        sleep $opt{delay};
        exit; 
    };

    alarm($opt{timeout});
    $ssh = "cat $opt{script} | " . $ssh
        if (exists $opt{script});
    $ssh = "cat $opt{copy} | " . $ssh
        if (exists $opt{copy});
    @pids = open(SSH, "$ssh $ssh_args $machine \'($exec_cmd) 2>&1\' 2>&1 |")
        or die "Cannot fork ssh: $!";

    open(LOGFILE, ">$opt{logpath}/$hostname")
        if (exists $opt{logpath});

    while (<SSH>)
    {
        chomp $_;
        $linecounter++;

        my $rline = ((exists $opt{suppress}) ? '' : "$hostname")
                  . ((exists $opt{linenums}) ? "\[$linecounter\]" : '')
                  . ((exists $opt{suppress}) ? '' : ': ')
                  . "$_\n";

        if ( (exists $opt{unbuffered}) and (not exists $opt{logpath}) )
        {
            print $rline;
        }
        elsif (exists $opt{logpath})
        {
            print LOGFILE "$_\n";
        }
        else
        {
            push(@output, $rline);
        }
    }

    print @output unless exists $opt{unbuffered};
    unless (exists $opt{quiet}) {
        print "." if exists $opt{printdots};
    }
    close(LOGFILE)
        if (exists $opt{logpath});

    close(SSH);
    alarm(0);

    sleep $opt{delay};
    exit;
}


sub reaper 
{
    my $pid;
    $SIG{CHLD} = sub { \&reaper };
    $pid = waitpid(-1, &WNOHANG);

    if (WIFEXITED($?)) 
    {
        $children--;
        delete $children{$pid};
    }
}


sub resolve_address
{
    my $address = shift; 
    my $hostname = gethostbyaddr(inet_aton($address), AF_INET);

    return $hostname if ($hostname) ;
    return $address;
}


sub build_hostlist
{
    my $syslist = shift;
    my %hostlist;

    open(SYSLIST, "<$opt{syslist}") 
        or die "Cannot open syslist: $!";

    foreach my $line (<SYSLIST>)
    {
        chomp $line;
        my @list = split(/\s+/,$line);

        foreach my $host (@list)
        {
            next if ($host =~ m/^\s*#/);

            if ($host =~ m/^(\d+\.\d+\.\d+\.\d+)/)
            {
                $hostlist{unreachable}{$host} = resolve_address($1);
            }
            else
            {
                $hostlist{unreachable}{$host} = $host;
            }
        }
    }
    close(SYSLIST);

    if (exists $opt{ping})
    {
        print 'Pinging ' . scalar(keys %{$hostlist{unreachable}})
              . " hosts ...\n" unless (exists $opt{quiet});

        my @fping_results;
        my @fping_targets = keys %{$hostlist{unreachable}};
        while (@fping_targets)
        {
            my @hosts = splice(@fping_targets, 0, 200);
            open(FPING, '-|', "$fping @hosts");
            push(@fping_results, <FPING>);
            close(FPING);
        }
    
        chomp(@fping_results);
        for my $result (@fping_results)
        {
            my ($host, $state) = ($result =~ /(\S+)\s+(.*)$/);
            if ($state eq 'is alive')
        {
                $hostlist{reachable}{$host} = $hostlist{unreachable}{$host};
                delete $hostlist{unreachable}{$host};
            }
        }
    }
    else
    {
        $hostlist{reachable} = $hostlist{unreachable};
        delete $hostlist{unreachable};
    }

    return \%hostlist;
}


sub check_ssh_keys
{
    my $keys = `$ssh_add -l 2>&1`;

    unless ($? == 0)
    {
        print STDERR "ERROR: No ssh keys loaded.\n";
        return 0;
    }
    return 1;
}


sub check_bash
{
    my $exec_cmd = shift;
    my @test_syntax = `$bash -nc \'$exec_cmd\' 2>&1`;

    if ($? == 0)
    {
        return 1;
    }
    else
    {    
        usage();
        print STDERR "ERROR: Invalid bash syntax:\n";
        for (@test_syntax)
        {
            chomp;
            print STDERR "\t$_\n";
        }

        return 0;
    }
}


sub print_list
{
    my @array_copy = @_;

    while (@array_copy)
    {
        my @line = splice(@array_copy, 0, 2);
        printf STDERR "%-35s" x @line . "\n", @line;
    }
}


sub check_tools
{
    my $tools = shift;
    my @path = split(/:/, $ENV{PATH});

    for my $tool (@$tools)
    {
        my $found = 0;
        for my $dir (@path)
        {
            if ( (-f "${dir}/${tool}") and (-x "${dir}/${tool}") )
            {
                $found = 1;
                last;
            }
        }

        unless ($found == 1)
        {
            print STDERR "ERROR: Could not find required tool $tool. "
                  . "Make sure it is available in your PATH.\n";
            return 0;
        }
    }
    return 1;
}


sub exit_error
{
    my $msg = shift;
    my $show_usage = shift;

    usage() if ($show_usage);
    print STDERR "$msg";
    exit 1;
}

sub usage
{
    my $extended = shift;
    my $progname = basename($0);

    print STDERR (<<EOF);

onall v${version} -- A parallelized system-administration tool.
Usage: onall [<options>] <command>
       onall [<options>] --script <script> <script_args>
       onall [<options>] --copy <file> <dest_path>

EOF
    if ($extended) 
    {
    print STDERR (<<EOF);
    -c, --copy <file> [<dest_path>]
                Copy <file> to destination as <dest_path>. If not
                specified, <dest_path> will be same as <file>.
                
    -d, --delay
                Delay N seconds between the fork of each ssh child
                process.  Default is 0.

    -h, --help
                Extended help (this screen).

    -l, --linenums
                Append line numbers to command generated output.
    
    -n, --no_syntax_check
                Skip the pre-flight syntax check of the command to run.
                Useful for issuing commands to targets whose shell is
                not bash.

    -o, --logpath
                Output results to files named by the hostname named
                by dir/hostname.

    -p, --ping
                Ping all hosts prior to attempting an ssh connection.
                Requires fping(8).

    -q, --quiet
                Semi-quiet mode.  Supresses sanity check information
                and assumes YES to all prompts.

    -r, --rate
                Run command on N machines in parallel.  This is the number
                of ssh connections open at a given time.  Default is 20.

    -e, --script <script> [<script_args>]
                Transfer/execute a script from the local system to each
                target system.  In this mode the <command> should be
                a path to a local script.  Any arguments which follow
                will be passed to the remote script.

    -A, --ssh_arguments
                Pass arguments directly to SSH.

    -k, --ssh_no_keys
                Do not require an ssh agent or keys.  Use with care.

    -Q, --ssh_quiet
                Run SSH in quiet mode.

    -s, --suppress
                Supress hostnames and line numbers from being prepended
                to result data.

    -f, --syslist
                Path to system list file.  If unspecified, the list will
                be read from STDIN.  Expected input format is hostname/IP,
                followed by newline.

    -t, --timeout
                Timeout in seconds before child ssh processes are killed.
                Be careful with this argument, as all running/pending
                commands on the machine in question will be terminated
                potentially leaving the system in an unknown state.
                Default is 300.

    -b, --unbuffered
                Turn OFF output buffering of result data.  By default,
                onall will display all output after remote command
                execution is complete, thereby maintaining proper order.
                When buffering is off, result data will be returned as
                each line is received.

    -u, --user
                Specifies the username ssh will use to log in to remote
                systems.  Default is root.

    -y, --yes
              Assume YES for all prompts.

    -x, --printdots
              Print a dot for each system to show progress.

Examples:

    onall -f cluster1.list -r50 -t30 'grep fls1 /etc/fstab | wc -l'
    cat cluster1web.list | onall -r20 '/etc/init.d/httpd-* stop'
    <tool that generates list of hosts based on options> | \
        onall -r 12 -t30 'df; vmstat'
    <tool that generates list of hosts based on options> | \
        | onall -e ~user/scripts/do_stuff --now
    <tool that generates list of hosts based on options> | \
        | onall -c ~bin/coolutil /usr/bin/coolutil
  
Notes:

    Unless invoked with the -n switch, command syntax is expected
    to be written in bash.    Syntax will be verified prior to forking
    any ssh children.  It is usually best to enclose the command in
    'single quotes', or insert shell escapes where nessessary.
    Unless invoked with the -k switch, SSH keys must be loaded via
    ssh-agent/ssh-add prior to running onall.

EOF
    }
}

